import matplotlib.pyplot as plt
import numpy as np


def draw_heatmap(value_map: np.ndarray, x_labels, y_labels, filepath: str, cmap=plt.cm.jet, show=False):
    assert len(x_labels) == value_map.shape[1]
    assert len(y_labels) == value_map.shape[0]

    fig, ax = plt.subplots(figsize=(10, 10))
    # 将元组分解为fig和ax两个变量
    mask = np.triu(np.ones_like(value_map))
    tril = value_map.copy()
    print(tril)
    tril[mask == 1] = np.nan
    im = ax.imshow(tril, cmap=cmap)
    # 显示图片

    ax.set_xticks(np.arange(len(x_labels)))
    # 设置x轴刻度间隔
    ax.set_yticks(np.arange(len(y_labels)))
    # 设置y轴刻度间隔
    ax.set_xticklabels(x_labels)
    # 设置x轴标签'''
    ax.set_yticklabels(y_labels)
    # 设置y轴标签'''

    plt.setp(ax.get_xticklabels(), rotation=90, ha="right",
             rotation_mode="anchor")
    # 设置标签 旋转45度 ha有三个选择：right,center,left（对齐方式）

    for i in range(len(y_labels)):
        for j in range(i):
            ax.text(j, i, f"{value_map[i, j]:.2f}", ha="center", va="center", color="grey")
    '''
    画图
    j,i:表示坐标值上的值
    value_map[i, j]表示内容
    ha有三个选择：right,center,left（对齐方式）
    va有四个选择：'top', 'bottom', 'center', 'baseline'（对齐方式）
    color:设置颜色
    '''
    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.ax.set_ylabel("Similarity", rotation=-90, va="bottom")

    # ax.set_title("Harvest of local x_labels (in tons/year)")
    # 设置题目
    fig.tight_layout()  # 自动调整子图参数,使之填充整个图像区域。
    plt.savefig(filepath)
    if show:
        plt.show()  # 图像展示
